package com.tianji.aigc.memory;

import com.tianji.common.utils.CollUtils;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.BoundListOperations;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.List;

/**
 * redis实现会话存储
 * @author zgp
 * @version V1.0
 * @date 2025-04-21 20:47
 */
public class RedisChatMemory implements ChatMemory {
    @Resource
    private StringRedisTemplate stringRedisTemplate;
    // 默认会话消息缓存key前缀
    private static final String DEFAULT_CHAT_KEY_PREFIX = "CHAT:";
    private final String chatKeyPrefix;

    public RedisChatMemory() {
        this.chatKeyPrefix = DEFAULT_CHAT_KEY_PREFIX;
    }

    public RedisChatMemory(String chatKeyPrefix) {
        this.chatKeyPrefix = chatKeyPrefix;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        if (CollUtils.isEmpty(messages)) {
            return;
        }
        // 将消息存储到redis中
        String key = getKey(conversationId);
        BoundListOperations<String, String> operations = stringRedisTemplate.boundListOps(key);
        // 将序列化后的消息存入redis的列表中（从右侧压栈，保证消息顺序）
        messages.forEach(message -> operations.rightPush(MessageUtil.toJson(message)));
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        if(lastN < 0){
            return List.of();
        }
        // 从reids中获取会话记录
        String key = getKey(conversationId);
        BoundListOperations<String, String> operations = stringRedisTemplate.boundListOps(key);
        // 得到消息json字符串列表
        List<String> messages = operations.range(0, lastN);
        if(CollUtils.isEmpty(messages)){
            return List.of();
        }
        // 反序列化消息
        return messages.stream()
                .map(MessageUtil::toMessage)
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        String key = getKey(conversationId);
        // 移除redis中的key
        stringRedisTemplate.delete(key);
    }

    /**
     * 获取历史会话记录key
     * @param conversationId 会话id
     * @return
     */
    private String getKey(String conversationId) {
        return chatKeyPrefix + conversationId;
    }

    /**
     * 根据对话ID优化对话记录，删除最后的2条消息，因为这2条消息是从路由智能体存储的，请求由后续的智能体处理
     * 为了确保历史消息的完整性，所以需要将中间转发的消息清理掉
     *
     * @param conversationId 对话的唯一标识符
     */
    public void optimization(String conversationId) {
        var redisKey = this.getKey(conversationId);
        var listOps = this.stringRedisTemplate.boundListOps(redisKey);
        // 从Redis列表右侧弹出2个元素
        listOps.rightPop(2);
    }
}
